"""
聚源基金数据分析 MCP 服务器
使用 FastMCP 构建，提供基金数据分析功能
支持工具分类以避免上下文瓶颈
"""

import logging
import os
import re
import sys
from typing import List, Optional, Dict, Any

from agno.agent import Agent
from agno.models.openai import OpenAILike
from pydantic_core.core_schema import StringSchema

from configs.basic_config import ALI_GENERATION_MODEL, AGENT_URL, ALI_DASHSCOPE_API_KEY
from configs.report_prompt import ppt_prompt, markdown_prompt, html_prompt

# 添加项目根目录到Python路径
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, project_root)

from mcp.server.fastmcp import FastMCP
from pydantic import BaseModel, Field
from tools.juyuan_tools.service.oracle_data_service import OracleDataService

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

try:
    #from mcp.server.fastmcp import FastMCP

    # 创建 FastMCP 实例 - 这是唯一的服务器实例
    #mcp = FastMCP("酒店管理 MCP 服务器")
    mcp = FastMCP("聚源基金数据分析 MCP 服务器", host='0.0.0.0', port=9001, stateless_http=True)
except ImportError:
    logger.error("❌ 无法导入 mcp，请安装: pip install mcp")
    raise

# 全局工具实例
oracle_data_service = None

logger.info("🚀 初始化聚源基金数据分析 MCP 服务器...")

# 初始化数据库
try:
    logger.info("💾 正在初始化Oracle数据库连接...")
    # 由于目录名包含连字符，需要特殊处理
    # tools_path = os.path.join(project_root, 'tools')
    # juyuan_tools_path = os.path.join(tools_path, 'juyuan_tools')
    #
    # # 添加到sys.path
    # if tools_path not in sys.path:
    #     sys.path.insert(0, tools_path)
    # if juyuan_tools_path not in sys.path:
    #     sys.path.insert(0, juyuan_tools_path)
    #
    # # 直接执行文件来获取类定义
    # service_file = os.path.join(juyuan_tools_path, 'service', 'oracle_data_service.py')
    # with open(service_file, 'r', encoding='utf-8') as f:
    #     service_code = f.read()
    #
    # # 创建一个模块并执行代码
    # import types
    # service_module = types.ModuleType('oracle_data_service')
    # exec(service_code, service_module.__dict__)
    #
    # OracleDataService = service_module.OracleDataService
    oracle_data_service = OracleDataService()
    logger.info("✅ Oracle数据库连接初始化成功.")
except Exception as e:
    logger.error(f"❌ 初始化失败: {e}")
    raise


# --- 参数模型定义 ---
class MCPGetFundBasicInfoParams(BaseModel):
    fund_code: Optional[str] = Field(None, description="基金代码")
    fund_name: Optional[str] = Field(None, description="基金名称")


class MCPMatchFundParams(BaseModel):
    fund_code: Optional[str] = Field(None, description="基金代码")
    fund_name: Optional[str] = Field(None, description="基金名称")


class MCPGetManagerInfoParams(BaseModel):
    innercode: int = Field(..., description="基金内部代码")


class MCPGetFundByManagerParams(BaseModel):
    manager_name: str = Field(..., description="基金经理名")


class MCPGetManagerViewParams(BaseModel):
    fund_code: str = Field(..., description="基金代码")
    report_year: int = Field(..., description="报告年份")
    report_quarter: int = Field(..., description="报告季度")


class MCPGetNetValueParams(BaseModel):
    fund_code: str = Field(..., description="基金代码")


class MCPGetFundScaleParams(BaseModel):
    innercode: int = Field(..., description="基金内部代码")


class MCPGetFundPerformanceParams(BaseModel):
    fund_code: Optional[str] = Field(None, description="基金代码")
    fund_name: Optional[str] = Field(None, description="基金名称")
    start_date: Optional[str] = Field(None, description="开始日期，格式为 YYYY-MM-DD")
    end_date: Optional[str] = Field(None, description="结束日期，格式为 YYYY-MM-DD")


class MCPGetTopTenStocksParams(BaseModel):
    fund_code: Optional[str] = Field(None, description="基金代码")
    fund_name: Optional[str] = Field(None, description="基金名称")
    start_date: Optional[str] = Field(None, description="开始日期，格式为 YYYY-MM-DD")
    end_date: Optional[str] = Field(None, description="结束日期，格式为 YYYY-MM-DD")


class MCPGetTopFiveBondsParams(BaseModel):
    fund_code: Optional[str] = Field(None, description="基金代码")
    fund_name: Optional[str] = Field(None, description="基金名称")
    start_date: Optional[str] = Field(None, description="开始日期，格式为 YYYY-MM-DD")
    end_date: Optional[str] = Field(None, description="结束日期，格式为 YYYY-MM-DD")


class MCPGetAllStocksParams(BaseModel):
    fund_code: Optional[str] = Field(None, description="基金代码")
    fund_name: Optional[str] = Field(None, description="基金名称")
    start_date: Optional[str] = Field(None, description="开始日期，格式为 YYYY-MM-DD")
    end_date: Optional[str] = Field(None, description="结束日期，格式为 YYYY-MM-DD")


class MCPGetAllBondsParams(BaseModel):
    fund_code: Optional[str] = Field(None, description="基金代码")
    fund_name: Optional[str] = Field(None, description="基金名称")
    start_date: Optional[str] = Field(None, description="开始日期，格式为 YYYY-MM-DD")
    end_date: Optional[str] = Field(None, description="结束日期，格式为 YYYY-MM-DD")


class MCPGetStocksChangeParams(BaseModel):
    fund_code: Optional[str] = Field(None, description="基金代码")
    fund_name: Optional[str] = Field(None, description="基金名称")
    start_date: Optional[str] = Field(None, description="开始日期，格式为 YYYY-MM-DD")
    end_date: Optional[str] = Field(None, description="结束日期，格式为 YYYY-MM-DD")


class MCPListCategoryToolsParams(BaseModel):
    """列出指定类别的工具"""
    category: str = Field(..., description="工具类别")


class MCPListAllCategoriesParams(BaseModel):
    """获取所有工具类别"""
    pass


class MCPHtmlReportParams(BaseModel):
    """生成HTML报告的参数"""
    data: List[Dict[str, Any]] = Field(..., description="要生成报告的数据")
    title: str = Field("基金数据分析报告", description="报告标题")
    description: str = Field("基于基金数据生成的分析报告", description="报告描述")


# --- 工具分类管理 ---
TOOL_CATEGORIES = {
    "fund_basic": {
        "name": "基金基础信息",
        "description": "提供基金基本信息查询功能",
        "tools": ["fund_basic.get_info", "fund_basic.match"]
    },
    "fund_manager": {
        "name": "基金经理相关",
        "description": "提供基金经理相关信息查询功能",
        "tools": ["fund_manager.get_info", "fund_manager.get_funds"]
    },
    "fund_performance": {
        "name": "基金业绩表现",
        "description": "提供基金业绩表现查询功能",
        "tools": ["fund_performance.get_data"]
    },
    "fund_holdings": {
        "name": "基金持仓信息",
        "description": "提供基金持仓信息查询功能",
        "tools": [
            "fund_holdings.get_top_stocks", 
            "fund_holdings.get_top_bonds",
            "fund_holdings.get_all_stocks",
            "fund_holdings.get_all_bonds",
            "fund_holdings.get_stocks_change"
        ]
    },
    "fund_net_value": {
        "name": "基金净值信息",
        "description": "提供基金净值信息查询功能",
        "tools": ["fund_net_value.get_data", "fund_net_value.get_scale"]
    },
    "report": {
        "name": "报告生成",
        "description": "提供数据可视化报告生成功能",
        "tools": ["report.generate_html"]
    }
}


# --- 工具注册 ---
@mcp.tool(
    name="fund_basic.get_info",
    description="获取基金基本信息"
)
async def get_fund_basic_info(
        fund_code: Optional[str] = None,
        fund_name: Optional[str] = None
) -> List[Dict[str, Any]]:
    """获取基金基本信息
    Args:
        fund_code: 基金代码（可选）
        fund_name: 基金全称（可选）
    """
    if not oracle_data_service:
        raise RuntimeError("Oracle数据服务未初始化")

    logger.info(f"🔍 获取基金基本信息: fund_code={fund_code}, fund_name={fund_name}")
    return oracle_data_service.get_fund_basic_info(fund_code, fund_name)


@mcp.tool(
    name="fund_basic.match",
    description="匹配基金代码,名称,内部代码"
)
async def match_fund(
        fund_code: Optional[str] = None,
        fund_name: Optional[str] = None
) -> List[Dict[str, Any]]:
    """匹配基金代码,名称，内部代码
    Args:
        fund_code: 基金代码（可选）
        fund_name: 基金全称（可选）
    """
    if not oracle_data_service:
        raise RuntimeError("Oracle数据服务未初始化")

    logger.info(f"🔍 匹配基金: fund_code={fund_code}, fund_name={fund_name}")
    return oracle_data_service.match_fund(fund_code, fund_name)


@mcp.tool(
    name="fund_manager.get_info",
    description="获取基金经理信息"
)
async def get_manager_info(innercode: int) -> List[Dict[str, Any]]:
    """获取基金经理信息
    Args:
        innercode: 基金内部代码（必选）
    """

    if not oracle_data_service:
        raise RuntimeError("Oracle数据服务未初始化")

    logger.info(f"👨‍💼 获取基金经理信息: innercode={innercode}")
    return oracle_data_service.get_manager_info(innercode)


@mcp.tool(
    name="fund_manager.get_funds",
    description="根据基金经理名查询基金"
)
async def get_fund_by_manager(manager_name: str) -> List[Dict[str, Any]]:
    """根据基金经理名查询基金
    Args:
        manager_name: 基金经理姓名（必选）
    """
    if not oracle_data_service:
        raise RuntimeError("Oracle数据服务未初始化")

    logger.info(f"🔍 根据基金经理名查询基金: manager_name={manager_name}")
    return oracle_data_service.get_fund_by_manager(manager_name)


@mcp.tool(
    name="fund_net_value.get_data",
    description="获取基金净值相关信息（包含单位净值和复权单位净值）"
)
async def get_net_value(
        fund_code: Optional[str] = None,
        fund_name: Optional[str] = None,
        fund_short_name: Optional[str] = None,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None
) -> List[Dict[str, Any]]:
    """获取基金净值相关信息（包含单位净值和复权单位净值）

    Args:
        fund_code: 基金代码（可选）
        fund_name: 基金全称（可选）
        fund_short_name: 基金简称（可选）
        start_date: 开始日期，格式为"YYYY-MM-DD"（可选）
        end_date: 结束日期，格式为"YYYY-MM-DD"（可选）
    """
    if not oracle_data_service:
        raise RuntimeError("Oracle数据服务未初始化")

    logger.info(
        f"💰 获取基金净值相关信息: fund_code={fund_code}, fund_name={fund_name}, fund_short_name={fund_short_name}, start_date={start_date}, end_date={end_date}")

    return oracle_data_service.get_net_value(
        fund_code=fund_code,
        fund_name=fund_name,
        fund_short_name=fund_short_name,
        start_date=start_date,
        end_date=end_date
    )


@mcp.tool(
    name="fund_net_value.get_scale",
    description="获取基金规模变化"
)
async def get_fund_scale(innercode: int) -> List[Dict[str, Any]]:
    """获取基金规模变化
    Args:
        innercode: 基金内部代码（必选）
    """
    if not oracle_data_service:
        raise RuntimeError("Oracle数据服务未初始化")

    logger.info(f"📊 获取基金规模变化: innercode={innercode}")
    return oracle_data_service.get_fund_scale(innercode)


@mcp.tool(
    name="fund_performance.get_data",
    description="获取基金业绩表现: 返回数据中-净值表现统计区间(StatisticRange)解释：1-一个月，2-两个月，3-季度，5-一季度(累计)，"
                "6-半年，11-三季度(累计)，12-年度，24-两年，36-三年，60-五年，99-成立至今，990-期内，993-截止时点，995-周，996-日，999-期末。"
)
async def get_fund_performance(
        fund_code: Optional[str] = None,
        fund_name: Optional[str] = None,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None
) -> List[Dict[str, Any]]:
    """获取基金业绩表现
    Args:
        fund_code: 基金代码（可选）
        fund_name: 基金全称（可选）
        start_date: 开始日期，格式为"YYYY-MM-DD"（可选）
        end_date: 结束日期，格式为"YYYY-MM-DD"（可选）
    """
    if not oracle_data_service:
        raise RuntimeError("Oracle数据服务未初始化")

    logger.info(f"📈 获取基金业绩表现: fund_code={fund_code}, fund_name={fund_name}, start_date={start_date}, end_date={end_date}")
    return oracle_data_service.get_fund_performance(
        fund_code=fund_code,
        fund_name=fund_name,
        start_date=start_date,
        end_date=end_date
    )


@mcp.tool(
    name="fund_holdings.get_top_stocks",
    description="获取基金的前十大重仓股"
)
async def get_top_ten_stocks(
        fund_code: Optional[str] = None,
        fund_name: Optional[str] = None,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None
) -> List[Dict[str, Any]]:
    """获取十大重仓股
    Args:
        fund_code: 基金代码（可选）
        fund_name: 基金全称（可选）
        start_date: 开始日期，格式为"YYYY-MM-DD"（可选）
        end_date: 结束日期，格式为"YYYY-MM-DD"（可选）
    """
    if not oracle_data_service:
        raise RuntimeError("Oracle数据服务未初始化")

    logger.info(f"🏢 获取十大重仓股: fund_code={fund_code}, fund_name={fund_name}, start_date={start_date}, end_date={end_date}")
    return oracle_data_service.get_top_ten_stocks(
        fund_code=fund_code,
        fund_name=fund_name,
        start_date=start_date,
        end_date=end_date
    )


@mcp.tool(
    name="fund_holdings.get_top_bonds",
    description="获取基金的前五大债券持仓"
)
async def get_top_five_bonds(
        fund_code: Optional[str] = None,
        fund_name: Optional[str] = None,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None
) -> List[Dict[str, Any]]:
    """获取前五大债券持仓
    Args:
        fund_code: 基金代码（可选）
        fund_name: 基金全称（可选）
        start_date: 开始日期，格式为"YYYY-MM-DD"（可选）
        end_date: 结束日期，格式为"YYYY-MM-DD"（可选）
    """
    if not oracle_data_service:
        raise RuntimeError("Oracle数据服务未初始化")

    logger.info(f"💵 获取前五大债券持仓: fund_code={fund_code}, fund_name={fund_name}, start_date={start_date}, end_date={end_date}")
    return oracle_data_service.get_top_five_bonds(
        fund_code=fund_code,
        fund_name=fund_name,
        start_date=start_date,
        end_date=end_date
    )


@mcp.tool(
    name="fund_holdings.get_all_stocks",
    description="获取全部股票持仓"
)
async def get_all_stocks(
        fund_code: Optional[str] = None,
        fund_name: Optional[str] = None,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None
) -> List[Dict[str, Any]]:
    """获取全部股票持仓
    Args:
        fund_code: 基金代码（可选）
        fund_name: 基金全称（可选）
        start_date: 开始日期，格式为"YYYY-MM-DD"（可选）
        end_date: 结束日期，格式为"YYYY-MM-DD"（可选）
    """
    if not oracle_data_service:
        raise RuntimeError("Oracle数据服务未初始化")

    logger.info(f"🏢 获取全部股票持仓: fund_code={fund_code}, fund_name={fund_name}, start_date={start_date}, end_date={end_date}")
    return oracle_data_service.get_all_stocks(
        fund_code=fund_code,
        fund_name=fund_name,
        start_date=start_date,
        end_date=end_date
    )


@mcp.tool(
    name="fund_holdings.get_all_bonds",
    description="获取全部债券持仓"
)
async def get_all_bonds(
        fund_code: Optional[str] = None,
        fund_name: Optional[str] = None,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None
) -> List[Dict[str, Any]]:
    """获取全部债券持仓
    Args:
        fund_code: 基金代码（可选）
        fund_name: 基金全称（可选）
        start_date: 开始日期，格式为"YYYY-MM-DD"（可选）
        end_date: 结束日期，格式为"YYYY-MM-DD"（可选）
    """
    if not oracle_data_service:
        raise RuntimeError("Oracle数据服务未初始化")

    logger.info(f"💵 获取全部债券持仓: fund_code={fund_code}, fund_name={fund_name}, start_date={start_date}, end_date={end_date}")
    return oracle_data_service.get_all_bonds(
        fund_code=fund_code,
        fund_name=fund_name,
        start_date=start_date,
        end_date=end_date
    )


@mcp.tool(
    name="fund_holdings.get_stocks_change",
    description="获取股票持仓变动; 返回结果中变动方向解释： 1-买入，2-卖出"
)
async def get_stocks_change(
        fund_code: Optional[str] = None,
        fund_name: Optional[str] = None,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None
) -> List[Dict[str, Any]]:
    """获取股票持仓变动
    Args:
        fund_code: 基金代码（可选）
        fund_name: 基金全称（可选）
        start_date: 开始日期，格式为"YYYY-MM-DD"（可选）
        end_date: 结束日期，格式为"YYYY-MM-DD"（可选）
    """
    if not oracle_data_service:
        raise RuntimeError("Oracle数据服务未初始化")

    logger.info(f"🔄 获取股票持仓变动: fund_code={fund_code}, fund_name={fund_name}, start_date={start_date}, end_date={end_date}")
    return oracle_data_service.get_stocks_change(
        fund_code=fund_code,
        fund_name=fund_name,
        start_date=start_date,
        end_date=end_date
    )


@mcp.tool(
    name="report.generate",
    description="""生成带有图表和数据可视化的分析报告。
    
    支持多种输出格式：
    1. ppt: PowerPoint演示文稿格式，适合汇报展示
    2. markdown: Markdown文档格式，适合文档编写
    3. html: HTML网页格式，适合网页展示
    
    参数说明:
    - data: 输入数据，可以是字典、列表或其他结构化数据
    - query: 用户查询，用于理解报告生成背景
    - title: 报告标题
    - description: 报告描述
    - report_type: 输出格式，默认为markdown
    """
)
async def generate_report(
        data: Optional[Any] = None,
        query: str = None,
        title: str = "数据分析报告",
        description: str = "基于数据生成的分析报告",
        report_type: str = "markdown"
) -> str:
    """生成带有图表和数据可视化的报告
    
    Args:
        data: 要生成报告的数据
        query: 用户提出的问题
        title: 报告标题
        description: 报告描述
        report_type: 报告类型，支持 "ppt"、"markdown"、"html"
    """

    # 定义报告生成器工厂
    report_factory = {
        "ppt": generate_ppt_report,
        "markdown": generate_markdown_report,
        "html": generate_html_report,
    }

    # 获取对应的报告生成器
    report_generator = report_factory.get(str(report_type).lower())
    
    if not report_generator:
        raise ValueError(f"不支持的报告类型: {report_type}")

    # 调用对应的报告生成器
    return await report_generator(data, query, title, description)


async def generate_ppt_report(data: Any, query: Optional[str], title: str, description: str) -> str:
    """生成PPT格式的报告"""
    # 构建提示词
    full_prompt = ppt_prompt.format(
        data=data,
        query=query,
        title=title,
        description=description
    )
    
    # 调用大模型生成报告
    agent = Agent(
        model=OpenAILike(id=ALI_GENERATION_MODEL, name=ALI_GENERATION_MODEL, api_key=ALI_DASHSCOPE_API_KEY,
                         base_url=AGENT_URL),
        instructions=full_prompt,
    )

    result = agent.run("依据指令，生成一份PPT格式的分析报告")
    
    patterns = [
        r'^```html\s*',  # 开头的 ```html
        r'\s*```$',  # 结尾的 ```
        r'^```\s*',  # 开头的 ```
        r'\s*```html$',  # 结尾的 ```html
    ]

    cleaned_content = result.content.strip() if result.content else ""

    for pattern in patterns:
        cleaned_content = re.sub(pattern, '', cleaned_content, flags=re.IGNORECASE)
    
    return cleaned_content


async def generate_markdown_report(data: Any, query: Optional[str], title: str, description: str) -> str:
    """生成Markdown格式的报告"""
    # 构建提示词
    full_prompt = markdown_prompt.format(
        data=data,
        query=query,
        title=title,
        description=description
    )
    
    # 调用大模型生成报告
    agent = Agent(
        model=OpenAILike(id=ALI_GENERATION_MODEL, name=ALI_GENERATION_MODEL, api_key=ALI_DASHSCOPE_API_KEY,
                         base_url=AGENT_URL),
        instructions=full_prompt,
    )

    result = agent.run("依据指令，生成一份Markdown格式的分析报告")
    
    return result.content.strip() if result.content else ""


async def generate_html_report(data: Any, query: Optional[str], title: str, description: str) -> str:
    """生成HTML格式的报告"""
    # 构建提示词
    full_prompt = html_prompt.format(
        data=data,
        query=query,
        title=title,
        description=description
    )
    
    # 调用大模型生成报告
    agent = Agent(
        model=OpenAILike(id=ALI_GENERATION_MODEL, name=ALI_GENERATION_MODEL, api_key=ALI_DASHSCOPE_API_KEY,
                         base_url=AGENT_URL),
        instructions=full_prompt,
    )

    result = agent.run("依据指令，生成一份HTML格式的分析报告")
    
    patterns = [
        r'^```html\s*',  # 开头的 ```html
        r'\s*```$',  # 结尾的 ```
        r'^```\s*',  # 开头的 ```
        r'\s*```html$',  # 结尾的 ```html
    ]

    cleaned_content = result.content.strip() if result.content else ""

    for pattern in patterns:
        cleaned_content = re.sub(pattern, '', cleaned_content, flags=re.IGNORECASE)
    
    return cleaned_content


async def debug_list_tools(mcp: FastMCP):
    """调试方法：打印所有工具信息"""
    tools = await mcp.list_tools()
    print("注册的工具:")
    print(tools)
    return tools

# --- 主程序入口 ---
if __name__ == "__main__":
    print("=" * 60)
    print("💼 聚源基金数据分析 MCP 服务器 ")
    print("-" * 60)
    print("支持工具分类，避免上下文瓶颈")
    print("工具类别:")
    for category, info in TOOL_CATEGORIES.items():
        print(f"  - {category}: {info['name']} ({len(info['tools'])}个工具)")

    # 在运行前打印工具信息
    import asyncio
    asyncio.run(debug_list_tools(mcp))

    # 使用 FastMCP 的内置运行方法
    mcp.run(transport="streamable-http")